# BSD3.0
# cili@163.com

import asyncio
import websockets
import json
import logging
import os
import base64
import re
from datetime import datetime
import hashlib

class WebSocketServer:
    def __init__(self, md_server, port=8000):
        self.connected = set()
        self.md_server = md_server
        self.port = port
        self.loop = None
        self.server = None
        self.should_exit = False
        self.is_running = False
        self.project_participants = {}
        with open('mdproject/projects.json','r',encoding='utf-8') as f:
            self.projects = json.load(f)
        self.MAX_IMAGES_PER_PROJECT = 20
        self.MAX_IMAGE_SIZE = 5 * 1024 * 1024  # 5MB

    def _get_image_md5(self, image_data):
        """计算图片数据的MD5值"""
        return hashlib.md5(image_data).hexdigest()

    def _clean_old_images(self, project_dir):
        """清理旧图片，保持项目图片数量在限制之内"""
        try:
            images = []
            for filename in os.listdir(project_dir):
                if filename.endswith('.jpg'):
                    file_path = os.path.join(project_dir, filename)
                    images.append((os.path.getmtime(file_path), file_path))
            
            # 如果图片数量超过限制，删除最旧的图片
            if len(images) >= self.MAX_IMAGES_PER_PROJECT:
                images.sort()  # 按时间戳排序
                for _, file_path in images[:(len(images) - self.MAX_IMAGES_PER_PROJECT + 1)]:
                    os.remove(file_path)
        except Exception as e:
            logging.error(f"Error cleaning old images: {e}")

    async def handle_image(self, project, image_data):
        """处理图片上传"""
        try:
            # 检查数据是否为二进制格式
            if isinstance(image_data, bytes):
                img_data = image_data
            else:
                # 如果是 base64 字符串，则进行解码
                header, encoded = image_data.split(",", 1)
                img_data = base64.b64decode(encoded)

            # 检查文件大小
            if len(img_data) > self.MAX_IMAGE_SIZE:
                raise ValueError("Image size exceeds maximum limit")

            # 计算MD5值作为文件名
            md5_hash = self._get_image_md5(img_data)
            filename = f"{md5_hash}.jpg"

            # 确保项目目录存在
            project_dir = os.path.join('mdproject', project, 'images')
            os.makedirs(project_dir, exist_ok=True)

            # 检查是否已存在相同MD5的图片
            filepath = os.path.join(project_dir, filename)
            if os.path.exists(filepath):
                return f"./mdproject/{project}/images/{filename}"

            # 检查并清理旧图片
            self._clean_old_images(project_dir)

            # 检查当前图片数量
            current_images = len([f for f in os.listdir(project_dir) 
                                if f.endswith('.jpg')])
            if current_images >= self.MAX_IMAGES_PER_PROJECT:
                raise ValueError("Maximum number of images reached for this project")

            # 保存新图片
            with open(filepath, 'wb') as f:
                f.write(img_data)

            # 通知所有相关客户端更新图片列表
            update_message = json.dumps({
                'type': 'imageListUpdate',
                'project': project
            })
            
            # 广播给项目中的所有客户端
            for conn in self.project_participants.get(project, set()):
                try:
                    if conn.state == websockets.protocol.State.OPEN:
                        await conn.send(update_message)
                except Exception as e:
                    logging.error(f"Error sending image list update: {e}")
                    continue

            return f"./mdproject/{project}/images/{filename}"

        except Exception as e:
            logging.error(f"Image processing error: {e}")
            return None

    async def delete_image(self, project, filename, content=None):
        """删除图片并通知所有客户端"""
        try:
            # 构建文件路径
            file_path = os.path.join('mdproject', project, 'images', filename)
            
            if not os.path.exists(file_path):
                return False, '文件不存在'
            
            # 删除文件
            os.remove(file_path)
            
            # 如果提供了新的文档内容，更新文档
            if content is not None:
                self.md_server.save_project_content(project, content)
            
            # 广播图片列表更新消息
            await self.broadcast_update(project, {
                'type': 'imageListUpdate',
                'project': project
            })
            
            # 如果更新了文档，广播文档更新消息
            if content is not None:
                await self.broadcast_update(project, {
                    'type': 'documentUpdate',
                    'project': project,
                    'content': content,
                    'timestamp': datetime.now().isoformat()
                })
            
            return True, '文件删除成功'
            
        except Exception as e:
            logging.error(f"Error deleting image: {e}")
            return False, f'删除文件失败: {str(e)}'

    async def broadcast_update(self, project, message):
        """广播更新消息给项目中的所有客户端"""
        if not isinstance(message, str):
            message = json.dumps(message)
        
        for conn in self.project_participants.get(project, set()):
            try:
                if conn.state == websockets.protocol.State.OPEN:
                    await conn.send(message)
            except Exception as e:
                logging.error(f"Error broadcasting update: {e}")
                continue

    def get_project_images(self, project):
        """获取项目的图片列表"""
        try:
            project_dir = os.path.join('mdproject', project, 'images')
            if not os.path.exists(project_dir):
                return []
            
            images = []
            for filename in os.listdir(project_dir):
                if filename.lower().endswith(('.jpg', '.jpeg', '.png', '.gif')):
                    file_path = os.path.join(project_dir, filename)
                    images.append({
                        'name': filename,
                        'timestamp': os.path.getmtime(file_path)
                    })
            
            # 按时间戳排序，最新的在前
            images.sort(key=lambda x: x['timestamp'], reverse=True)
            return images
        except Exception as e:
            logging.error(f"Error getting project images: {e}")
            return []

    async def handler(self, websocket):
        with open('mdproject/projects.json','r',encoding='utf-8') as f:
            self.projects = json.load(f)
        self.connected.add(websocket)
        logging.info(f"Client connected: {websocket.remote_address}")
        try:
            async for message in websocket:
                data = json.loads(message)
                message_type = data['type']
                print(message_type)
                if message_type == 'getImageList':
                    project = data.get('project')
                    if project in self.projects:
                        images = self.get_project_images(project)
                        await websocket.send(json.dumps({
                            'type': 'imageList',
                            'project': project,
                            'images': images
                        }))

                elif message_type == 'join':
                    project = data['project']
                    if project not in self.project_participants:
                        self.project_participants[project] = set()

                    if len(self.project_participants[project]) >= 5 and websocket not in self.project_participants[project]:
                        await websocket.send(json.dumps({
                            'type': 'joinResult',
                            'success': False,
                            'message': '项目参与者已达到上限(5人)'
                        }))
                        continue

                    if self.md_server.verify_project(project, data['password']):
                        self.project_participants[project].add(websocket)
                        content = self.md_server.get_project_content(project)
                        images = self.get_project_images(project)
                        
                        await websocket.send(json.dumps({
                            'type': 'joinResult',
                            'success': True,
                            'project': project,
                            'content': content
                        }))
                        
                        # 紧接着发送图片列表
                        await websocket.send(json.dumps({
                            'type': 'imageList',
                            'project': project,
                            'images': images
                        }))
                    else:
                        await websocket.send(json.dumps({
                            'type': 'joinResult',
                            'success': False,
                            'message': '密码错误'
                        }))

                elif message_type == 'update':
                    content = data['content']
                    project = data['project']
                    timestamp = data.get('timestamp', 0)
                    
                    # 立即保存和广播更新
                    self.md_server.save_project_content(project, content)
                    update_message = json.dumps({
                        'type': 'documentUpdate',
                        'project': project,
                        'content': content,
                        'timestamp': timestamp
                    })
                    
                    # 广播给其他客户端
                    for conn in self.project_participants.get(project, set()):
                        try:
                            if conn != websocket and conn.state == websockets.protocol.State.OPEN:
                                await conn.send(update_message)
                        except Exception as e:
                            logging.error(f"Error sending update: {e}")
                            continue

                elif message_type == 'upload':
                    logging.info("Handling upload request")
                    project = data.get('project')
                    image_data = data.get('data')

                    if project not in self.projects:
                        await websocket.send(json.dumps({
                            'type': 'uploadResult',
                            'success': False,
                            'message': '项目不存在'
                        }))
                        continue

                    try:
                        image_url = await self.handle_image(project, image_data)
                        if image_url:
                            await websocket.send(json.dumps({
                                'type': 'uploadResult',
                                'success': True,
                                'url': image_url
                            }))
                        else:
                            await websocket.send(json.dumps({
                                'type': 'uploadResult',
                                'success': False,
                                'message': '图片处理失败'
                            }))
                    except ValueError as ve:
                        await websocket.send(json.dumps({
                            'type': 'uploadResult',
                            'success': False,
                            'message': str(ve)
                        }))
                    except Exception as e:
                        logging.error(f"Upload error: {str(e)}")
                        await websocket.send(json.dumps({
                            'type': 'uploadResult',
                            'success': False,
                            'message': str(e)
                        }))

                elif message_type == 'deleteImage':
                    project = data.get('project')
                    filename = data.get('filename')
                    content = data.get('content')  # 可选的文档内容

                    if project not in self.projects:
                        await websocket.send(json.dumps({
                            'type': 'deleteResult',
                            'success': False,
                            'message': '项目不存在'
                        }))
                        continue

                    success, message = await self.delete_image(project, filename, content)
                    await websocket.send(json.dumps({
                        'type': 'deleteResult',
                        'success': success,
                        'message': message
                    }))

        except Exception as e:
            logging.error(f"Error: {e}")
        finally:
            for project in self.project_participants:
                if websocket in self.project_participants[project]:
                    self.project_participants[project].remove(websocket)
            self.connected.remove(websocket)
            logging.info(f"Client disconnected: {websocket.remote_address}")

    async def main(self):
        try:
            self.server = await websockets.serve(
                self.handler,
                "0.0.0.0",
                self.port,
                reuse_address=True
            )
            self.is_running = True
            await asyncio.Future()
        except Exception as e:
            logging.error(f"Server start error: {e}")
            self.is_running = False
            raise e

    async def stop_server(self):
        """停止WebSocket服务器"""
        try:
            # 关闭所有连接
            connected_copy = list(self.connected)
            for ws in connected_copy:
                try:
                    await ws.close()
                except:
                    pass
            self.connected.clear()

            # 关闭服务器
            if self.server:
                self.server.close()
                try:
                    await asyncio.wait_for(self.server.wait_closed(), timeout=2.0)
                except asyncio.TimeoutError:
                    pass  # 如果等待超时，继续执行
        except Exception as e:
            logging.error(f"Error stopping server: {e}")

    def start(self):
        self.should_exit = False
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)
        try:
            self.loop.run_until_complete(self.main())
        finally:
            self.loop.close()

    def stop(self):
        if self.loop and self.is_running:
            async def cleanup():
                await self.stop_server()
                self.is_running = False

            if self.loop.is_running():
                asyncio.run_coroutine_threadsafe(cleanup(), self.loop)
            
            # 等待服务器完全停止
            import time
            timeout = 0
            while self.is_running and timeout < 3:  # 最多等待3秒
                time.sleep(0.1)
                timeout += 0.1

            if self.loop:
                self.loop.stop()
            self.loop = None
            self.server = None

    async def handle_upload(self, websocket, data):
        try:
            logging.info("Starting upload handling")  # 调试信息
            project = data.get('project')
            image_data = data.get('data')

            logging.info(f"Upload for project: {project}")  # 调试信息

            # 验证项目
            if project not in self.projects:
                logging.error(f"Project not found: {project}")  # 调试信息
                await websocket.send(json.dumps({
                    'type': 'uploadResult',
                    'success': False,
                    'message': '不存在'
                }))
                return

            # ... rest of the code ...

        except Exception as e:
            logging.error(f"Upload error: {str(e)}")
            logging.exception(e)  # 打印完整的错误堆栈
            await websocket.send(json.dumps({
                'type': 'uploadResult',
                'success': False,
                'message': str(e)
            }))
